# -*- coding:utf-8 -*-
# @Time    : 2023/5/11 13:10
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : websocket_server.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# 通过websocket与前端进行通信
# 1. 接收前端的消息
# 2. 将消息传递给agent
# 3. 将agent的回复传递给前端
import os
import re
import sys
import threading
import uuid
from datetime import datetime

from agent.llm_agent_parent import lock
from bot.insurance_planner_gpt import company_config
from bot.insurance_planner_gpt.config.conf import profile_
from bot.insurance_planner_gpt.context_local_util import context_data
from bot.insurance_planner_gpt.utils import date_util
from bot.insurance_planner_gpt.utils.conversation_process_util import ConversationProcessUtil
from bot.insurance_planner_gpt.utils.planner_msg_buffer_util import planner_msg_buffer_util
from bot.insurance_planner_gpt.utils.seat_planner_msg_channel import SeatPlannerMsgChannel
from common.common_exception import CommonException

from websockets.exceptions import ConnectionClosedOK, ConnectionClosed
import common.constants
from bot.insurance_planner_gpt.agent.user_suggestion import UserSuggestion
from common import constants

import traceback
from common.message import ConversationMessage, MessageContent
import asyncio
import websockets
from common.log import logger
import ssl
from bot.insurance_planner_gpt.planner import PlannerChat

from db.redis.redis_service import RedisMemory
from util_tool.conversation_util import ConversationUtil
import settings
import openai

certfile = './common/ssl/socket.chatlingxi.com.pem'
keyfile = './common/ssl/socket.chatlingxi.com.key'

# 创建SSL上下文
ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
ssl_context.load_cert_chain(certfile=certfile, keyfile=keyfile)
redis_memory = RedisMemory()

pro_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
sys.path.append(pro_path)
from data_generate import utils


class SocketProcessSeatChannel(SeatPlannerMsgChannel):

    def __init__(self):
        super().__init__()

    async def handle_seat_channel_data(self, clients=None):
        """
        规划师服务订阅redis队列下的seat_channel 接收坐席转发的消息 把坐席修改后的内容通过socket传输回用户
        """
        pubsub = self.aio_redis.pubsub()
        await pubsub.subscribe(self.seat_channel)
        try:
            async for ori_message in pubsub.listen():
                if ori_message['type'] == 'message':
                    message = ori_message['data'].decode('utf-8')
                    context_message_all = ConversationMessage.from_json(message)
                    user_session_id = context_message_all.session_id
                    messages = context_message_all.messages
                    client_type = context_message_all.client_type
                    client_id = context_message_all.client_id
                    # 初始化时客户端发送的消息
                    last_message = messages[-1]
                    role = last_message.get('role', None)
                    message_id = last_message.get('message_id', str(uuid.uuid4()).replace('-', ''))
                    message_id_process = redis_memory.hget("planner_channel:message_id", message_id)
                    if not message_id_process:
                        redis_memory.hset("planner_channel:message_id", message_id, 1)
                        redis_memory.expireat("planner_channel:message_id", date_util.get_next_day_start())
                        last_content_data = last_message['content'][0].data
                        logger.info(
                            f"会话id:{user_session_id},从{self.seat_channel}中获取到一条用户消息:{context_message_all.to_json()}")
                        if clients and user_session_id in clients:
                            websocket = clients[user_session_id]
                            message = ConversationMessage.gen_msg(constants.MESSAGE_TYPE_TEXT, client_id,
                                                                  client_type, last_content_data,
                                                                  user_session_id, user_session_id,
                                                                  code=constants.MESSAGE_CODE_END,
                                                                  role=role, message_id=message_id)
                            await websocket.send(message.to_json())
        except Exception as e:
            logger.error(f"处理{self.seat_channel}渠道消息异常: {str(e)}")
        finally:
            await pubsub.close()
            await pubsub.unsubscribe(self.seat_channel)


class WebsocketServer:

    def __init__(self):
        self.cu = ConversationUtil()
        self.clients = {}
        redis_memory.delete("active_sessions:socket")
        self.channel = SocketProcessSeatChannel()

    def post_process_code(self, code):
        sep = "\n```"
        if sep in code:
            blocks = code.split(sep)
            if len(blocks) % 2 == 1:
                for i in range(1, len(blocks), 2):
                    blocks[i] = blocks[i].replace("\\_", "_")
            code = sep.join(blocks)
        return code

    async def send_message(self, planner_chat, context, session_id, client_type, user_name, client_id):
        loop = asyncio.get_event_loop()
        message = ""
        # 消息id 流式消息和结束消息使用同一个消息id
        message_id = str(uuid.uuid4()).replace('-', '')
        context_data.set({"session_id": session_id, "message_id": message_id})
        # stream_iter, other_messages = await planner_chat.async_reply(context, session_id, client_type, is_stream=True)
        is_stream = False
        stream_iter, other_messages = await planner_chat.async_reply(context, session_id, client_type,
                                                                     is_stream=is_stream)
        # 查询当前会话是否被接管
        takeover_client_id = redis_memory.get("agent_takeover:" + session_id)

        if other_messages:
            await self.send_additional_message(client_id, client_type, loop, other_messages, planner_chat, session_id,
                                               user_name, message_id)
        try:
            if is_stream:
                async for data in stream_iter:
                    if ("choices" in data):
                        message_dict = data['choices'][0]["delta"]
                        if "content" in message_dict:
                            message = message + message_dict["content"]
                    else:
                        message = data["text"]
                        message = message.replace("�", "")
                        message = self.post_process_code(message)
                    #  如果当前会话被接管，则不发送消息给用户
                    if not takeover_client_id:
                        await planner_chat.send_msg(constants.MESSAGE_TYPE_TEXT, client_type, session_id, message,
                                                    constants.MESSAGE_CODE_SUCCESS, message_id, client_id)
            else:
                message = stream_iter
            logger.info(f"会话id:{session_id} ,话术生成输出的结果是: {message}")
            # 流式输出的prompt和结果保存
            prompt = redis_memory.get(f"planner:Conversation:prompt:{session_id}:{message_id}")
            if prompt:
                logger.info(
                    f"当前会话:{session_id},当前任务：Conversation , the stream prompt is:\n {prompt}\n{message}")
                loop.run_in_executor(None, self.save_train_data, prompt, message,
                                     "Conversation", "llm", "stream",
                                     session_id + ":" + message_id)
                redis_memory.delete(f"planner:Conversation:prompt:{session_id}:{message_id}")
        except websockets.exceptions.ConnectionClosed as e:
            logger.error(f"socket关闭,会话id:{session_id}：" + traceback.format_exc())
            raise CommonException(9999, message)
        # 消息为空处理
        if not message:
            raise Exception
        # 刷新下redis的值 发送最终结果前 如果坐席退出接管直接发给用户
        takeover_client_id = redis_memory.get("agent_takeover:" + session_id)
        if not takeover_client_id:
            # 如果发送过程中坐席退出接管或者无人接管直接把最终消息发回给用户
            await planner_chat.send_msg(constants.MESSAGE_TYPE_TEXT, client_type, session_id, message,
                                        constants.MESSAGE_CODE_END, message_id, client_id)
        # 发送规划师响应给坐席工作台
        await self.send_planner_result_to_seat(client_id, client_type, message, message_id, planner_chat, session_id)
        # 存储大模型的最新对话
        if user_name and user_name[:4] != "LXKJ":
            logger.info(f"会话id:{session_id}，正常响应，存储规划师的响应：{message}")
            loop.run_in_executor(None, self.cu.store_conversation_index, session_id,
                                 "origin_" + message_id,
                                 constants.MESSAGE_TYPE_TEXT, message,
                                 constants.ROLE_ASSISTANT, client_type)

        # 查询当前会话是否被接管
        takeover_client_id = redis_memory.get("agent_takeover:" + session_id)
        # 发送用户建议
        if not takeover_client_id:
            await self.gen_user_suggestion(planner_chat, context, message, session_id,
                                           client_type, client_id)

        return message_id, message

    async def send_planner_result_to_seat(self, client_id, client_type, message, message_id, planner_chat, session_id):
        try:
            msg = ConversationMessage.gen_msg(constants.MESSAGE_TYPE_EDIT_TEXT, client_id, client_type, message,
                                              session_id, role=constants.ROLE_ASSISTANT, message_id=message_id,
                                              action="takeover_response")
            await self.channel.publish_msg(self.channel.planner_channel, msg)

        except Exception:
            logger.info(f"当前会话id:{session_id},给所有正在连接的坐席发送用户最新的对话异常,{traceback.format_exc()}")

    async def send_additional_message(self, client_id, client_type, loop, other_messages, planner_chat, session_id,
                                      user_name, message_id):
        """ 其他类型消息的处理"""
        try:
            # 发送 图片信息给用户
            for item in other_messages:
                logger.info(f"会话id:{session_id}，发送给坐席工作台的额外消息为：{item.data}")
                if isinstance(item, MessageContent):
                    if item.msg_type == constants.MESSAGE_TYPE_IMG:
                        await planner_chat.send_msg(constants.MESSAGE_TYPE_IMG, client_type, session_id, "▌",
                                                    constants.MESSAGE_CODE_SUCCESS, None, client_id)
                        await planner_chat.send_msg(constants.MESSAGE_TYPE_IMG, client_type, session_id, item.data,
                                                    constants.MESSAGE_CODE_SUCCESS, message_id, client_id)
                        if user_name and user_name[:4] != "LXKJ":
                            loop.run_in_executor(None, self.cu.store_conversation_index, session_id, message_id,
                                                 constants.MESSAGE_TYPE_IMG,
                                                 item.data,
                                                 constants.ROLE_ASSISTANT, client_type)
                    elif item.msg_type == constants.MESSAGE_TYPE_AGENT_SOLUTION:
                        # 发送给坐席的方案信息
                        msg = ConversationMessage.gen_msg(constants.MESSAGE_TYPE_AGENT_SOLUTION, client_id, client_type,
                                                          item, session_id, role=constants.ROLE_ASSISTANT)
                        await self.channel.publish_msg(self.channel.planner_channel, msg)
                        if user_name and user_name[
                                         :4] != "LXKJ" and item.msg_type == constants.MESSAGE_TYPE_AGENT_SOLUTION:
                            loop.run_in_executor(None, self.cu.store_conversation_index, session_id,
                                                 'origin_' + message_id,
                                                 constants.MESSAGE_TYPE_AGENT_SOLUTION,
                                                 item.data,
                                                 constants.ROLE_ASSISTANT, client_type)

        except Exception:
            pass

    async def echo(self, websocket):
        session_id = ""
        client_type = ""
        user_name = ""
        client_id = ""
        try:
            loop = asyncio.get_event_loop()
            # async 会被同步IO阻塞，要么IO改成异步，要么使用线程池
            async for message in websocket:
                planner_chat = None
                try:
                    # 解析消息
                    context_message_all = ConversationMessage.from_json(message)
                    session_id = context_message_all.session_id
                    # 存入协程上下文  后面取用
                    context_data.set({"session_id": session_id, "message_id": 'message_id'})
                    client_type = context_message_all.client_type
                    config = company_config.get_config_from_client_type(client_type)
                    planner_role = config.get("planner_role", "")
                    company_name = config.get("company_name", "")
                    company_business = config.get("company_business", "")
                    company_values = config.get("company_values", "")
                    # 实例化小助对象 每一轮对话都实例化
                    planner_chat = PlannerChat(websocket, planner_role, company_name, company_business, company_values)
                    messages = context_message_all.messages
                    code = context_message_all.code
                    user_msg_receive_time = datetime.now()
                    client_id = context_message_all.client_id
                    # client_id默认为session_id
                    if not client_id:
                        client_id = context_message_all.session_id
                    user_name = context_message_all.user_name
                    try:
                        # 记录用户的活跃会话
                        self.clients[client_id] = websocket
                        redis_memory.sadd('active_sessions:socket', session_id)
                        # 启动与坐席工作台的通道
                        asyncio.create_task(self.channel.handle_seat_channel_data(self.clients))
                    except:
                        logger.error(f"会话id:{session_id},连接坐席工作台通道异常!")
                    context = ConversationMessage.format_context(messages)
                    conversation = ConversationProcessUtil.format_conversation_history(context,
                                                                                       assistant_role_name=planner_role)
                    # 初始化时客户端发送的消息
                    if messages:
                        last_message = messages[-1]
                        role = last_message.get('role', None)
                        message_id = last_message.get('message_id')
                        if not message_id:
                            message_id = str(uuid.uuid4()).replace('-', '')
                        last_content_data = last_message['content'][0].data
                        logger.info(f"会话id:{session_id},收到一条最新的用户对话===={last_content_data}")
                        msg_type = last_message['content'][0].msg_type
                        # 查询当前是否被处理了 多次相同消息的处理

                        # buffer = planner_msg_buffer_util.get_top_from_buffer(session_id)
                        # if buffer and last_content_data:
                        #     if buffer.get('msg') == last_content_data:
                        #         logger.info(f"当前会话:[{session_id}]  当前处理的消息和缓冲区的消息不一致, 放弃处理该消息")
                        #         return
                        # # 把消息放到缓冲队列
                        # if last_content_data:
                        #     planner_msg_buffer_util.push_msg2buffer(session_id, last_content_data)
                        await self.store_user_session(conversation, loop, session_id, user_name)
                        # 存储用户最新的对话
                        if ((last_content_data and not code) or (
                                code and code != constants.MESSAGE_CODE_RETRY)) and user_name and user_name[
                                                                                                  :4] != "LXKJ":
                            logger.info(f"会话id:{session_id},存储最新的用户对话：{last_content_data}")
                            loop.run_in_executor(None, self.cu.store_conversation_index, session_id, message_id,
                                                 msg_type,
                                                 last_content_data,
                                                 role, client_type, user_msg_receive_time)
                        try:
                            # 给所有正在连接的坐席发送用户最新的对话 丢到队列
                            msg = ConversationMessage.gen_msg(msg_type, client_id, client_type, last_content_data,
                                                              session_id, role=constants.ROLE_USER)
                            await self.channel.publish_msg(self.channel.planner_channel, msg)

                        except Exception:
                            logger.info(f"当前会话id:{session_id},给所有正在连接的坐席发送用户最新的对话异常,{traceback.format_exc()}")
                        # 没有响应前先发送 ▌
                    await planner_chat.send_msg(constants.MESSAGE_TYPE_TEXT, client_type, session_id, "▌",
                                                constants.MESSAGE_CODE_SUCCESS, None, client_id)
                    await self.send_message(planner_chat, context, session_id, client_type,
                                            user_name, client_id)
                    # planner_msg_buffer_util.pop_msg2buffer(session_id)
                except ConnectionClosed:
                    # 捕获异常，并将异常信息发送给客户端
                    logger.error("连接关闭:" + traceback.format_exc())
                except CommonException as e:
                    if user_name and user_name[:4] != "LXKJ" and e.code == 9999 and e.msg:
                        logger.info(
                            f"流式输出中断，code={e.code},msg={e.msg},session={session_id},异常为：{traceback.format_exc()}")
                        loop.run_in_executor(None, self.cu.store_conversation_index, session_id,
                                             str(uuid.uuid4()).replace('-', ''), common.constants.MESSAGE_TYPE_TEXT,
                                             e.msg,
                                             common.constants.ROLE_ASSISTANT, client_type)
                        await planner_chat.send_msg(constants.MESSAGE_TYPE_TEXT, client_type, session_id, "连接中断",
                                                    constants.MESSAGE_CODE_TIMEOUT,
                                                    str(uuid.uuid4()).replace('-', ''), client_id)
                except openai._exceptions.BadRequestError:
                    logger.error(traceback.format_exc())
                    # token超出限制
                    logger.info(f"token超出限制。。会话id:{session_id},异常为{traceback.format_exc()}")
                    await planner_chat.send_msg(constants.MESSAGE_TYPE_TEXT, client_type, session_id,
                                                "字符输入到限制了,请开启新的对话～",
                                                constants.MESSAGE_CODE_TOKEN_LIMIT,
                                                str(uuid.uuid4()).replace('-', ''), client_id)
                except TimeoutError:
                    # 超时
                    logger.info(f"服务超时,session_id:{session_id},异常为{traceback.format_exc()}")
                    await planner_chat.send_msg(constants.MESSAGE_TYPE_TEXT, client_type, session_id, "系统超时，重试中···",
                                                constants.MESSAGE_CODE_TIMEOUT, str(uuid.uuid4()).replace('-', ''),
                                                client_id)
                except Exception:
                    # 捕获异常，并将异常信息发送给客户端
                    logger.info(f"服务异常,session_id:{session_id},异常为{traceback.format_exc()}")
                    await planner_chat.send_msg(constants.MESSAGE_TYPE_TEXT, client_type, session_id,
                                                "系统异常",
                                                constants.MESSAGE_CODE_TIMEOUT, str(uuid.uuid4()).replace('-', ''),
                                                client_id)
        except ConnectionClosed:
            # 连接正常关闭的处理
            logger.error("连接关闭-------" + traceback.format_exc())
            redis_memory.delete("socket_planner:dialogues_process:" + session_id)
        finally:
            # 用户连接关闭后，从活跃连接集合中移除
            redis_memory.srem("active_sessions:socket", session_id)
            self.remove_socket(websocket)
            logger.info(f"会话：{session_id}结束，清理连接")

    async def store_user_session(self, conversation, loop, session_id, user_name):
        """存储用户账号和会话的关系"""
        try:
            if user_name and user_name[:4] != "LXKJ":
                logger.info(f"会话id:{session_id},存储用户会话")
                # 存储用户<->会话关系
                loop.run_in_executor(None, self.cu.store_conversation_user, session_id, user_name,
                                     conversation)
        except Exception:
            logger.error(f"会话id:{session_id},存储会话异常：{traceback.format_exc()}")

    async def gen_user_suggestion(self, planner_chat, context, message, session_id, client_type, client_id):
        # extracted_strings_arr = ["我想要了解保险规划", "我想要了解个人保障", "我想咨询下医疗险"]
        extracted_strings_arr = [""]
        try:
            conversation = ConversationProcessUtil.format_conversation_history(context,
                                                                               assistant_role_name=planner_chat.planner_role)
            user_suggestion_obj = UserSuggestion(conversation_history=conversation)
            # user_suggestion = await user_suggestion_obj.achat_auto_llm(type='gpt')
            user_suggestion = ''
            logger.info(f"会话id:{session_id},生成用户回复建议，建议是：{user_suggestion}")
            pattern = r'\[(.*?)\]'
            user_suggestion = user_suggestion.replace(" ", "").replace(''"，"'', ''","'')
            matches = re.search(pattern, user_suggestion)
            if matches:
                extracted_string = matches.group(1)
                extracted_strings_arr = extracted_string.split('","')
        except Exception:
            logger.error(f"生成用户建议异常：{traceback.format_exc()}")
        finally:
            await planner_chat.send_msg(constants.MESSAGE_TYPE_SUGGESTION, client_type, session_id,
                                        ','.join(
                                            [str(item).replace(',', "，").replace('"', '') for item in
                                             extracted_strings_arr]),
                                        constants.MESSAGE_CODE_SUCCESS, str(uuid.uuid4()).replace('-', ''), client_id)

    def start(self, port):
        if settings.ONLINE:
            start_server = websockets.serve(self.echo, "0.0.0.0", port, ssl=ssl_context)
        else:
            start_server = websockets.serve(self.echo, "0.0.0.0", port, )
        asyncio.get_event_loop().run_until_complete(start_server)
        asyncio.get_event_loop().run_forever()

    def remove_socket(self, websocket):
        try:
            # 当客户端断开连接时，从映射中移除用户
            clients_to_remove = []
            for client_id, client in self.clients.items():
                if client == websocket:
                    clients_to_remove.append(client_id)
            for client_id in clients_to_remove:
                del self.clients[client_id]
        except Exception:
            pass

    def save_train_data(self, prompt, result, classname, file_name="llm", model_name="gpt-4", id=str(uuid.uuid1())):
        try:
            with lock:
                project_dir = os.path.dirname(os.path.dirname(__file__))
                # 带上日期的文件名yyyy-mm-dd
                now_date = datetime.now().strftime("%Y-%m-%d")
                now_datetime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

                path = os.path.join(project_dir, "internal_server/train_data/" + file_name + "/" + now_date + ".json")
                if not os.path.exists(path):
                    new_data = []
                else:
                    new_data = utils.jload(path)
                instruct_dict_new = {}
                instruct_dict_new["id"] = id
                instruct_dict_new["time"] = now_datetime
                instruct_dict_new["task"] = classname
                instruct_dict_new["model_name"] = model_name
                instruct_dict_new["conversations"] = []

                conversation_dict = {}
                conversation_dict["from"] = "human"
                conversation_dict["value"] = prompt
                instruct_dict_new["conversations"].append(conversation_dict)

                conversation_dict = {}
                conversation_dict["from"] = "gpt"
                conversation_dict["value"] = result
                instruct_dict_new["conversations"].append(conversation_dict)
                new_data.append(instruct_dict_new)
                utils.jdump(new_data, path)
        except Exception as ee:
            logger.error("save data error: {}".format(ee))


if __name__ == '__main__':
    websocket_server = WebsocketServer()
    port = profile_['websocket_server']['port']
    logger.info(f"********** websocket server started on port:{port} **********")
    websocket_server.start(port=port)
